{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rF2x3qooyBTI"
   },
   "source": [
    "# Deep Convolutional Generative Adversarial Network\n",
    "\n",
    "**Learning Objectives**\n",
    "\n",
    "- Build a GAN architecture (consisting of a generator and discriminator) in Keras \n",
    "- Define the loss for the generator and discriminator\n",
    "- Define a training step for the GAN using `tf.GradientTape()` and `@tf.function`\n",
    "- Train the GAN on the MNIST dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ITZuApL56Mny"
   },
   "source": [
    "## Introduction\n",
    "This notebook demonstrates how to build and train a [Generative Adversarial Network](https://arxiv.org/abs/1406.2661) (GAN) to generate images of handwritten digits using a [Deep Convolutional Generative Adversarial Network](https://arxiv.org/pdf/1511.06434.pdf) (DCGAN).\n",
    "\n",
    "GANs consist of two models which are trained simultaneously through an adversarial process. A *generator* (\"the artist\") learns to create images that look real, while a *discriminator* (\"the art critic\") learns to tell real images apart from fakes.\n",
    "\n",
    "![A diagram of a generator and discriminator](../assets/gan1.png)\n",
    "\n",
    "During training, the *generator* progressively becomes better at creating images that look real, while the *discriminator* becomes better at recognizing fake images. The process reaches equilibrium when the *discriminator* can no longer distinguish real images from fakes.\n",
    "\n",
    "![A second diagram of a generator and discriminator](../assets/gan2.png)\n",
    "\n",
    "In this notebook we'll build a GAN to generate MNIST digits. This notebook demonstrates this process on the MNIST dataset. The following animation shows a series of images produced by the *generator* as it was trained for 50 epochs. The images begin as random noise, and increasingly resemble hand written digits over time.\n",
    "\n",
    "![sample output](https://tensorflow.org/images/gan/dcgan.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "e1_Y75QXJS6h"
   },
   "source": [
    "## Import TensorFlow and other libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "J5oue0oqCkZZ"
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import, division, print_function, unicode_literals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "g5RstiiB8V-z"
   },
   "outputs": [],
   "source": [
    "try:\n",
    "  %tensorflow_version 2.x\n",
    "except Exception:\n",
    "  pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WZKbyU2-AiY-"
   },
   "outputs": [],
   "source": [
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wx-zNbLqB4K8"
   },
   "outputs": [],
   "source": [
    "tf.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YzTlj4YdCip_"
   },
   "outputs": [],
   "source": [
    "# To generate GIFs\n",
    "!python3 -m pip install -q imageio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "YfIk2es3hJEd"
   },
   "outputs": [],
   "source": [
    "import glob\n",
    "import imageio\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import os\n",
    "import PIL\n",
    "from tensorflow.keras import layers\n",
    "import time\n",
    "\n",
    "from IPython import display"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "iYn4MdZnKCey"
   },
   "source": [
    "## Load and prepare the dataset\n",
    "\n",
    "For this notebook, we will use the MNIST dataset to train the generator and the discriminator. The generator will generate handwritten digits resembling the MNIST data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "a4fYMGxGhrna"
   },
   "outputs": [],
   "source": [
    "(train_images, train_labels), (_, _) = tf.keras.datasets.mnist.load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NFC2ghIdiZYE"
   },
   "outputs": [],
   "source": [
    "train_images = train_images.reshape(train_images.shape[0], 28, 28, 1).astype('float32')\n",
    "train_images = (train_images - 127.5) / 127.5 # Normalize the images to [-1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "S4PIDhoDLbsZ"
   },
   "outputs": [],
   "source": [
    "BUFFER_SIZE = 60000\n",
    "BATCH_SIZE = 256"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, we define our input pipeline using `tf.data`. The pipeline below reads in `train_images` as tensor slices and then shuffles and batches the examples for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "-yKCCQOoJ7cn"
   },
   "outputs": [],
   "source": [
    "# Batch and shuffle the data\n",
    "train_dataset = tf.data.Dataset.from_tensor_slices(train_images)\n",
    "train_dataset = train_dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "THY-sZMiQ4UV"
   },
   "source": [
    "## Create the generator and discriminator models\n",
    "\n",
    "Both our generator and discriminator models will be defined using the [Keras Sequential API](https://www.tensorflow.org/guide/keras#sequential_model)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-tEyxE-GMC48"
   },
   "source": [
    "### The Generator\n",
    "\n",
    "The generator uses [tf.keras.layers.Conv2DTranspose](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Conv2DTranspose) (upsampling) layers to produce an image from a seed (random noise). We will start with a `Dense` layer that takes this seed as input, then upsample several times until you reach the desired image size of 28x28x1."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.** Complete the code below to create the generator model. Start with a dense layer that takes as input random noise. We will create random noise using `tf.random.normal([1, 100])`. Use `tf.keras.layers.Conv2DTranspose` over multiple layers to upsample the random noise from dimension 100 to ultimately dimension 28x28x1 (the shape of our original MNIST digits).\n",
    "\n",
    "Hint: Experiment with using `BatchNormalization` or different activation functions like `LeakyReLU`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "6bpTcDqoLWjY"
   },
   "outputs": [],
   "source": [
    "#TODO 1\n",
    "def make_generator_model():\n",
    "    model = tf.keras.Sequential()\n",
    "    \n",
    "    # TODO: Your code goes here.\n",
    "    \n",
    "    assert model.output_shape == (None, 28, 28, 1)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "GyWgG09LCSJl"
   },
   "source": [
    "Let's use the (as yet untrained) generator to create an image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gl7jcC7TdPTG"
   },
   "outputs": [],
   "source": [
    "generator = make_generator_model()\n",
    "\n",
    "noise = tf.random.normal([1, 100])\n",
    "generated_image = generator(noise, training=False)\n",
    "\n",
    "plt.imshow(generated_image[0, :, :, 0], cmap='gray')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "D0IKnaCtg6WE"
   },
   "source": [
    "### The Discriminator\n",
    "\n",
    "Next, we will build the discriminator. The discriminator is a CNN-based image classifier. It should take in an image of shape 28x28x1 and return a single classification indicating if that image is real or not."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.** Complete the code below to create the CNN-based discriminator model. Your model should be binary classifier which takes as input a tensor of shape 28x28x1. Experiment with different stacks of convolutions, activation functions, and/or dropout."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "dw2tPLmk2pEP"
   },
   "outputs": [],
   "source": [
    "#TODO 1.\n",
    "def make_discriminator_model():\n",
    "    model = tf.keras.Sequential()\n",
    "    \n",
    "    # TODO: Your code goes here.\n",
    "    \n",
    "    assert model.output_shape == (None, 1)\n",
    "\n",
    "    return model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `.summary()` we can have a high-level summary of the generator and discriminator models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_generator_model().summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "make_discriminator_model().summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "QhPneagzCaQv"
   },
   "source": [
    "Let's use the (as yet untrained) discriminator to classify the generated images as real or fake. The model will be trained to output positive values for real images, and negative values for fake images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "gDkA05NE6QMs"
   },
   "outputs": [],
   "source": [
    "discriminator = make_discriminator_model()\n",
    "decision = discriminator(generated_image)\n",
    "print(decision)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "0FMYgY_mPfTi"
   },
   "source": [
    "## Define the loss and optimizers\n",
    "\n",
    "Next, we will define the loss functions and optimizers for both the generator and discriminaotr models. Both the generator and discriminator will use the `BinaryCrossentropy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "psQfmXxYKU3X"
   },
   "outputs": [],
   "source": [
    "# This method returns a helper function to compute cross entropy loss\n",
    "cross_entropy = tf.keras.losses.BinaryCrossentropy(from_logits=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "PKY_iPSPNWoj"
   },
   "source": [
    "### Discriminator loss\n",
    "\n",
    "The method below quantifies how well the discriminator is able to distinguish real images from fakes. \n",
    "\n",
    "Recall, when training the discriminator (i.e. holding the generator fixed) the loss function has two parts: the loss when sampling from the real data and the loss when sampling from the fake data. The function below compares the discriminator's predictions on real images to an array of 1s, and the discriminator's predictions on fake (generated) images to an array of 0s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.**\n",
    "Complete the code in the method below. The `real_loss` should return the cross-entropy for the discriminator's predictions on real images and the `fake_loss` should return the cross-entropy for the discriminator's predictions on fake images."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "wkMNfBWlT-PV"
   },
   "outputs": [],
   "source": [
    "#TODO 2\n",
    "def discriminator_loss(real_output, fake_output):\n",
    "    real_loss = # TODO: Your code goes here.\n",
    "    fake_loss = # TODO: Your code goes here.\n",
    "    total_loss = # TODO: Your code goes here.\n",
    "    return total_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Jd-3GCUEiKtv"
   },
   "source": [
    "### Generator loss\n",
    "The generator's loss quantifies how well it was able to trick the discriminator. Intuitively, if the generator is performing well, the discriminator will classify the fake images as real (or 1). Here, we will compare the discriminators decisions on the generated images to an array of 1s."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.**\n",
    "Complete the code to return the cross-entropy loss of the generator's output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "90BIcCKcDMxz"
   },
   "outputs": [],
   "source": [
    "#TODO 2\n",
    "def generator_loss(fake_output):\n",
    "    return # Your code goes here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "MgIc7i0th_Iu"
   },
   "source": [
    "### Optimizers for the geneerator and discriminator\n",
    "\n",
    "Note that we must define two separete optimizers for the discriminator and the generator optimizers since we will train two networks separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "iWCn_PVdEJZ7"
   },
   "outputs": [],
   "source": [
    "generator_optimizer = tf.keras.optimizers.Adam(1e-4)\n",
    "discriminator_optimizer = tf.keras.optimizers.Adam(1e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "mWtinsGDPJlV"
   },
   "source": [
    "### Save checkpoints\n",
    "This notebook also demonstrates how to save and restore models, which can be helpful in case a long running training task is interrupted."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "CA1w-7s2POEy"
   },
   "outputs": [],
   "source": [
    "checkpoint_dir = \"./gan_training_checkpoints\"\n",
    "checkpoint_prefix = os.path.join(checkpoint_dir, \"ckpt\")\n",
    "checkpoint = tf.train.Checkpoint(generator_optimizer=generator_optimizer,\n",
    "                                 discriminator_optimizer=discriminator_optimizer,\n",
    "                                 generator=generator,\n",
    "                                 discriminator=discriminator)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "Rw1fkAczTQYh"
   },
   "source": [
    "## Define the training loop\n",
    "\n",
    "Next, we define the training loop for training our GAN. Below we set up global variables for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "NS2GWywBbAWo"
   },
   "outputs": [],
   "source": [
    "EPOCHS = 50\n",
    "noise_dim = 100\n",
    "num_examples_to_generate = 16\n",
    "\n",
    "# We will reuse this seed overtime (so it's easier)\n",
    "# to visualize progress in the animated GIF)\n",
    "seed = tf.random.normal([num_examples_to_generate, noise_dim])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "jylSonrqSWfi"
   },
   "source": [
    "The training loop begins with generator receiving a random seed as input. That seed is used to produce an image. The discriminator is then used to classify real images (drawn from the training set) and fakes images (produced by the generator). The loss is calculated for each of these models, and the gradients are used to update the generator and discriminator."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Exercise.**\n",
    "Complete the code below to define the training loop for our GAN. Notice the use of `tf.function` below. This annotation causes the function `train_step` to be \"compiled\". The `train_step` function takes as input a batch of images. In the rest of the function,\n",
    "- `generated_images` is created using the `generator` function with `noise` as input\n",
    "- apply the discriminator model to the `images` and `generated_images` to create the `real_output` and `fake_output` (resp.)\n",
    "- define the `gen_loss` and `disc_loss` using the methods you defined above.\n",
    "- compute the gradients of the generator and the discriminator using `gen_tape` and `disc_tape` (resp.)\n",
    "\n",
    "Lastly, we use the `.apply_gradients` method to make a gradient step for the `generator_optimizer` and `discriminator_optimizer`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "3t5ibNo05jCB"
   },
   "outputs": [],
   "source": [
    "# TODO 3\n",
    "@tf.function\n",
    "def train_step(images):\n",
    "    noise = tf.random.normal([BATCH_SIZE, noise_dim])\n",
    "\n",
    "    with tf.GradientTape() as gen_tape, tf.GradientTape() as disc_tape:\n",
    "        generated_images = # TODO: Your code goes here.\n",
    "\n",
    "        real_output = discriminator(images, training=True)\n",
    "        fake_output = discriminator(generated_images, training=True)\n",
    "\n",
    "        gen_loss = generator_loss(fake_output)\n",
    "        disc_loss = discriminator_loss(real_output, fake_output)\n",
    "\n",
    "    gradients_of_generator = gen_tape.gradient(\n",
    "        gen_loss, generator.trainable_variables)\n",
    "    gradients_of_discriminator = disc_tape.gradient(\n",
    "        disc_loss, discriminator.trainable_variables)\n",
    "\n",
    "    generator_optimizer.apply_gradients(\n",
    "        zip(gradients_of_generator, generator.trainable_variables))\n",
    "    discriminator_optimizer.apply_gradients(\n",
    "        zip(gradients_of_discriminator, discriminator.trainable_variables))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We use the `train_step` function above to define training of our GAN. Note here, the `train` function takes as argument the `tf.data` dataset and the number of epochs for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "2M7LmLtGEMQJ"
   },
   "outputs": [],
   "source": [
    "def train(dataset, epochs):\n",
    "    for epoch in range(epochs):\n",
    "        start = time.time()\n",
    "\n",
    "        for image_batch in dataset:\n",
    "            train_step(image_batch)\n",
    "\n",
    "        # Produce images for the GIF as we go\n",
    "        display.clear_output(wait=True)\n",
    "        generate_and_save_images(generator,\n",
    "                                 epoch + 1,\n",
    "                                 seed)\n",
    "\n",
    "        # Save the model every 15 epochs\n",
    "        if (epoch + 1) % 15 == 0:\n",
    "            checkpoint.save(file_prefix=checkpoint_prefix)\n",
    "\n",
    "        print ('Time for epoch {} is {} sec'.format(\n",
    "            epoch + 1, time.time()-start))\n",
    "\n",
    "    # Generate after the final epoch\n",
    "    display.clear_output(wait=True)\n",
    "    generate_and_save_images(generator,\n",
    "                             epochs,\n",
    "                             seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "2aFF7Hk3XdeW"
   },
   "source": [
    "**Generate and save images.**\n",
    "We'll use a small helper function to generate images and save them. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "RmdVsmvhPxyy"
   },
   "outputs": [],
   "source": [
    "def generate_and_save_images(model, epoch, test_input):\n",
    "    # Notice `training` is set to False.\n",
    "    # This is so all layers run in inference mode (batchnorm).\n",
    "    predictions = model(test_input, training=False)\n",
    "\n",
    "    fig = plt.figure(figsize=(4,4))\n",
    "\n",
    "    for i in range(predictions.shape[0]):\n",
    "            plt.subplot(4, 4, i+1)\n",
    "            plt.imshow(predictions[i, :, :, 0] * 127.5 + 127.5,\n",
    "                       cmap='gray')\n",
    "            plt.axis('off')\n",
    "\n",
    "    plt.savefig('./gan_images/image_at_epoch_{:04d}.png'.format(epoch))\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dZrd4CdjR-Fp"
   },
   "source": [
    "## Train the model\n",
    "Call the `train()` method defined above to train the generator and discriminator simultaneously. Note, training GANs can be tricky. It's important that the generator and discriminator do not overpower each other (e.g., that they train at a similar rate).\n",
    "\n",
    "At the beginning of the training, the generated images look like random noise. As training progresses, the generated digits will look increasingly real. After about 50 epochs, they resemble MNIST digits. This may take about one ot two minutes / epoch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Ly3UN0SLLY2l"
   },
   "outputs": [],
   "source": [
    "# TODO 4\n",
    "# TODO: Your code goes here."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "rfM4YcPVPkNO"
   },
   "source": [
    "Restore the latest checkpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "XhXsd0srPo8c"
   },
   "outputs": [],
   "source": [
    "checkpoint.restore(tf.train.latest_checkpoint(checkpoint_dir))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "P4M_vIbUi7c0"
   },
   "source": [
    "## Create a GIF\n",
    "\n",
    "Lastly, we'll create a gif that shows the progression of our produced images through training. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "WfO5wCdclHGL"
   },
   "outputs": [],
   "source": [
    "# Display a single image using the epoch number\n",
    "def display_image(epoch_no):\n",
    "    return PIL.Image.open('./gan_images/image_at_epoch_{:04d}.png'.format(epoch_no))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "5x3q9_Oe5q0A"
   },
   "outputs": [],
   "source": [
    "display_image(EPOCHS)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "NywiH3nL8guF"
   },
   "source": [
    "Use `imageio` to create an animated gif using the images saved during training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "IGKQgENQ8lEI"
   },
   "outputs": [],
   "source": [
    "anim_file = 'dcgan.gif'\n",
    "\n",
    "with imageio.get_writer(anim_file, mode='I') as writer:\n",
    "    filenames = glob.glob('./gan_images/image*.png')\n",
    "    filenames = sorted(filenames)\n",
    "    last = -1\n",
    "    for i,filename in enumerate(filenames):\n",
    "        frame = 2*(i**0.5)\n",
    "        if round(frame) > round(last):\n",
    "            last = frame\n",
    "        else:\n",
    "            continue\n",
    "        image = imageio.imread(filename)\n",
    "        writer.append_data(image)\n",
    "    image = imageio.imread(filename)\n",
    "    writer.append_data(image)\n",
    "\n",
    "import IPython\n",
    "if IPython.version_info > (6,2,0,''):\n",
    "    display.Image(filename=anim_file)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "k6qC-SbjK0yW"
   },
   "source": [
    "## Next steps\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "xjjkT9KAK6H7"
   },
   "source": [
    "This tutorial has shown the complete code necessary to write and train a GAN. As a next step, you might like to experiment with a different dataset, for example the Large-scale Celeb Faces Attributes (CelebA) dataset [available on Kaggle](https://www.kaggle.com/jessicali9530/celeba-dataset). To learn more about GANs we recommend the [NIPS 2016 Tutorial: Generative Adversarial Networks](https://arxiv.org/abs/1701.00160).\n",
    "\n",
    "Copyright 2020 Google Inc. Licensed under the Apache License, Version 2.0 (the \"License\"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "dcgan.ipynb",
   "private_outputs": true,
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
